package edu.iastate.cs228.hw5;
import java.util.AbstractSet;
import java.util.ArrayList;
import java.util.Iterator;
import java.util.List;
import java.util.NoSuchElementException;

/**
 * Binary search tree implementation of the Set interface.  The contains() and
 * remove() methods of AbstractSet are overridden to search the tree without
 * using the iterator.  Instances of this class always maintain node counts; that
 * is, the count() method of the BSTNode interface is implemented to be O(1).
 * If constructed with the isSelfBalancing flag true, instances of this tree are 
 * self-balancing: whenever an add(), remove(), or Iterator.remove() operation 
 * causes any node to become unbalanced, a rebalance operation is automatically
 * performed at the highest unbalanced node.
 */
public class BalancedBSTSet<E extends Comparable<? super E>> extends AbstractSet<E>
{
  /**
   * Root of this tree.
   */
  protected Node root;
  
  /**
   * Number of elements in this tree.
   */
  protected int size;
  
  /**
   * Tells if tree is self-balancing or not
   */
  protected boolean selfBalancing;
  
  /**
   * Represents the alpha balancing factor
   */
  protected int top;
  
  /**
   * Represents the alpha balancing factor
   */
  protected int bottom;
  
  /**
   * Node type for this implementation.
   */
  protected class Node implements BSTNode<E>
  {    
    public Node left;
    public Node right;
    public Node parent;
    public E data;
    public int subCount;
    
    public Node(E key, Node parent)
    {
      this.data = key;
      this.parent = parent;
      subCount = 1;
    }
    
    @Override
    public BSTNode<E> left()
    {
      return left;
    }

    @Override
    public BSTNode<E> right()
    {
      return right;
    }

    @Override
    public int count()
    {
      return subCount;
    }

    @Override
    public E data()
    {
      return data;
    }

    @Override
    public BSTNode<E> parent()
    {
      return parent;
    }
    
    @Override
    public String toString()
    {
      return data.toString();
    }
  } 
  /**
   * Performs an inorder traversal on this subtree
   * @param root
   * 	- the root of the given subtree
   * @param ordered
   * 	- List in which in order nodes are stored
   * @return
   * 	- The complete list of ordered nodes
   */
  protected List<BSTNode<E>> inorderTraversal (BSTNode<E> root, List<BSTNode<E>> ordered)
  {
	  if (root != null)
	  {
		  inorderTraversal(root.left(), ordered);
		  ordered.add(root);
		  inorderTraversal(root.right(), ordered);
	  }
	  
	  return ordered;
  }

  /**
   * Constructs an empty binary search tree.  This tree will maintain
   * node counts but will not automatically perform rebalance operations.
   */
  public BalancedBSTSet()
  {
	  top = 2;
	  bottom = 3;
	  selfBalancing = false;
  }

  /**
   * Constructs an empty binary search tree. If the isSelfBalancing
   * flag is true, the tree will be self-balancing: if so, whenever an add(), 
   * remove(), or Iterator.remove() operation causes any node to become 
   * unbalanced, a rebalance operation is automatically performed at the 
   * highest unbalanced node.  The default value alpha = 2/3 is used for
   * the balance condition.  Maintains node counts whether or not 
   * isSelfBalancing is true.
   * 
   * @param isSelfBalancing true if this binary search tree is 
   *   to be self-balancing, false otherwise
   */
  public BalancedBSTSet(boolean isSelfBalancing)
  {
	  selfBalancing = isSelfBalancing;
	  top = 2;
	  bottom = 3;
  }
  

  /**
   * Constructs an empty binary search tree. If the isSelfBalancing
   * flag is true, the tree will be self-balancing: if so, whenever an add(), 
   * remove(), or Iterator.remove() operation causes any node to become 
   * unbalanced, a rebalance operation is automatically performed at the 
   * highest unbalanced node.  The given alpha = top/bottom is used for
   * the balance condition.  Maintains node counts whether or not 
   * isSelfBalancing is true.
   * 
   * @param isSelfBalancing true if this binary search tree is 
   *   to be self-balancing, false otherwise
   * @param top numerator of the fraction alpha
   * @param bottom denominator of the fraction alpha
   * @throws IllegalArgumentException if top / bottom is less than 1/2 
   */
  public BalancedBSTSet(boolean isSelfBalancing, int top, int bottom)
  {
	  selfBalancing = isSelfBalancing;
	  this.top = top;
	  this.bottom = bottom;
	  double alpha = (double) top / (double) bottom;
	  if (alpha < 0.5)
		  throw new IllegalArgumentException();
  }  

  /**
   * Returns a read-only view of the root node of this tree.
   * @return root node of this tree
   */
  public BSTNode<E> root()
  {
    return root;
  }

  /**
   * This helper method is used to check if the tree is in need of rebalancing
   * @param node
   * 	- The root of the subtree to be balanced
   */
  private void checkBalance(Node node)
  {	
	  if (node==null)
		  return;
	  
	  int leftSize;
	  int rightSize;

	  Node highestUnbalanced = null;
	  Node temp = node;
	  while (temp != null)
	  {
		  leftSize = (temp.left==null) ? 0 : temp.left.subCount;
		  rightSize = (temp.right==null) ? 0 : temp.right.subCount;
		  
		  if (!(leftSize * bottom <= temp.subCount * top) || !(rightSize * bottom <= temp.subCount * top))
		  {
			  highestUnbalanced = temp;
		  }
		  temp = temp.parent;
	  }
	  if (highestUnbalanced != null)
	  {
		  rebalance(highestUnbalanced);
	  }
	  
  }
  
  /**
   * Performs a rebalance operation on the given subtree.  This operation
   * does not create or destroy any nodes and does not change the size of
   * this tree.
   * @param bstNode root of the subtree to be rebalanced.
   * @throws IllegalArgumentException if the given node is not part of this tree.
   * @throws ClassCastException if the given node cannot be cast to the 
   *   correct concrete node type for this tree.
   */
  public void rebalance(BSTNode<E> bstNode)
  {
	  
	  if (findEntry(bstNode.data()) == null)
		  throw new IllegalArgumentException();
	  if (!(bstNode instanceof BSTNode<?>))
		  throw new ClassCastException();
	  

		  List<BSTNode<E>> ordered = new ArrayList<BSTNode<E>>();
		  ordered = inorderTraversal(bstNode, ordered); // array list containing nodes in sequential order
		  int mid = (ordered.size()-1)/2;
		  Node midNode = (Node) ordered.get(mid);
		  Node initNode = (Node) bstNode;
		  if (bstNode.parent() != null)
		  {
			  if (initNode.parent().right() == initNode)
			  {
				  midNode.parent = (Node) bstNode.parent();
				  initNode.parent.right = midNode;
				  rebalanceRec(ordered, midNode, 0, ordered.size()-1);
			  }
			  else
			  {
				  midNode.parent = (Node) bstNode.parent();
				  initNode.parent.left = midNode;
				  rebalanceRec(ordered, midNode, 0, ordered.size()-1);
			  }
		  }
		  else
		  {
			  root = midNode;
			  root.parent = null;
			  rebalanceRec(ordered, midNode, 0, ordered.size()-1);
		  }
		  
  }
  		
  /**
   * Recursive portion of the rebalance function
   */
  		private void rebalanceRec (List<BSTNode<E>> ordered, Node parent, int first, int last)
  		{
  			int mid = (last+first)/2;
  			Node curMid = (Node) ordered.get(mid);
  			if (first == last) // last node
  			{
  				
  				curMid.left = null;
  				curMid.right = null;
  				curMid.parent = parent;
  				curMid.subCount = 1;
  				return;
  			}
  			else if (first == (last-1)) // last two nodes
  			{
  				Node left = (Node) ordered.get(first);
  				Node right = (Node) ordered.get(last);
  				left.left = null;
  				left.right = right;
  				left.parent = parent;
  				left.subCount = 2;
  				right.left = null;
  				right.right = null;
  				right.parent = left;
  				right.subCount = 1;
  				return;
  			}
  			else // normal split
  			{
  				int leftMid = ((mid-1)+first)/2;
  				int rightMid = (last+(mid+1))/2;
  				
  				curMid.left = (Node) ordered.get(leftMid);
  				curMid.right = (Node) ordered.get(rightMid);
  				curMid.left.parent = curMid;
  				curMid.right.parent = curMid;
  				
  				curMid.subCount = (last-first)+1;
  				
  				rebalanceRec(ordered, curMid, first, mid-1);
  				rebalanceRec(ordered, curMid, mid+1, last);
  			}
  		}
  
  
  @Override
  public boolean contains(Object obj)
  {
    // This cast may cause comparator to throw ClassCastException at runtime,
    // which is the expected behavior
    E key = (E) obj;
    return findEntry(key) != null;
  }
  
  @Override
  public boolean add(E key)
  {
    if (root == null)
    {
      root = new Node(key, null);
      ++size;
      return true;
    }
    
    Node current = root;
    while (true)
    {
      int comp = current.data.compareTo(key);
      if (comp == 0)
      {
        // key is already in the tree
        return false;
      }
      else if (comp > 0)
      {
        if (current.left != null)
        {
          current = current.left;
        }
        else
        {
          current.left = new Node(key, current);
          
          // update subtree size counter
          Node temp = current;
          while (temp != null)
          {
        	  temp.subCount++;
        	  temp = temp.parent;
          }
          if (selfBalancing)
        	  checkBalance(current.left);

          
          ++size;
          return true;
        }
      }
      else
      {
        if (current.right != null)
        {
          current = current.right;
        }
        else
        {
          current.right = new Node(key, current);
          
          
          // update subtree size counter
          Node temp = current;
          while (temp != null)
          {
        	  temp.subCount++;
        	  temp = temp.parent;
          }
          if (selfBalancing)
        	  checkBalance(current.right);
          ++size;
          return true;
        }
      }
    }   
  }
  
  @Override
  public boolean remove(Object obj)
  {
    // This cast may cause comparator to throw ClassCastException at runtime,
    // which is the expected behavior
    E key = (E) obj;
    Node n = findEntry(key);
    if (n == null)
    {
      return false;
    }
    unlinkNode(n);
    return true;
  }
  
  /**
   * Returns the node containing key, or null if the key is not
   * found in the tree.
   * @param key
   * @return the node containing key, or null if not found
   */
  protected Node findEntry(E key)
  {
    Node current = root;
    while (current != null)
    {
      int comp = current.data.compareTo(key);
      if (comp == 0)
      {
        return current;
      }
      else if (comp > 0)
      {
        current = current.left;
      }
      else
      {
        current = current.right;
      }
    }   
    return null;
  }
  

  
  /**
   * Returns the successor of the given node.
   * @param n
   * @return the successor of the given node in this tree, 
   *   or null if there is no successor
   */
  protected Node successor(Node n)
  {
    if (n == null)
    {
      return null;
    }
    else if (n.right != null)
    {
      // leftmost entry in right subtree
      Node current = n.right;
      while (current.left != null)
      {
        current = current.left;
      }
      return current;
    }
    else 
    {
      // we need to go up the tree to the closest ancestor that is
      // a left child; its parent must be the successor
      Node current = n.parent;
      Node child = n;
      while (current != null && current.right == child)
      {
        child = current;
        current = current.parent;
      }
      // either current is null, or child is left child of current
      return current;
    }
  }
  
  /**
   * Removes the given node, preserving the binary search
   * tree property of the tree.
   * @param n node to be removed
   */
  protected void unlinkNode(Node n)
  {
    // first deal with the two-child case; copy
    // data from successor up to n, and then delete successor 
    // node instead of given node n
    if (n.left != null && n.right != null)
    {
      Node s = successor(n);
      n.data = s.data;
      n = s; // causes s to be deleted in code below
    }
    
    // n has at most one child
    Node replacement = null;    
    if (n.left != null)
    {
      replacement = n.left;
    }
    else if (n.right != null)
    {
      replacement = n.right;
    }
    
    // link replacement into tree in place of node n 
    // (replacement may be null)
    if (n.parent == null)
    {
      root = replacement;
    }
    else
    {
      if (n == n.parent.left)
      {
        n.parent.left = replacement;
      }
      else
      {
        n.parent.right = replacement;
      }
    }
    
    if (replacement != null)
    {
      replacement.parent = n.parent;
    }
    
    // update subtree counts
    Node temp = n.parent;
    while (temp != null)
    {
    	temp.subCount--;
    	temp = temp.parent;
    }
    
    if (selfBalancing)
    	checkBalance(n.parent);
    --size;
  }
  
  @Override
  public Iterator<E> iterator()
  {
    return new BSTIterator();
  }

  @Override
  public int size()
  {
    return size;
  }
 
  /**
   * Returns a representation of this tree as a multi-line string.
   * The tree is drawn with the root at the left and children are
   * shown top-to-bottom.  Leaves are marked with a "-" and non-leaves
   * are marked with a "+".
   */
  @Override
  public String toString()
  {
    StringBuilder sb = new StringBuilder();
    toStringRec(root, sb, 0);
    return sb.toString();
  }
  
  /**
   * Preorder traversal of the tree that builds a string representation
   * in the given StringBuilder.
   * @param n root of subtree to be traversed
   * @param sb StringBuilder in which to create a string representation
   * @param depth depth of the given node in the tree
   */
  private void toStringRec(Node n, StringBuilder sb, int depth)
  {
    for (int i = 0; i < depth; ++i)
    {
      sb.append("  ");
    }
    
    if (n == null)
    {
      sb.append("-\n");
      return;
    }
    
    if (n.left != null || n.right != null)
    {
      sb.append("+ ");
    }
    else
    {
      sb.append("- ");
    }
    sb.append(n.data.toString() + " (" + n.count() + ")");
    sb.append("\n");
    if (n.left != null || n.right != null)
    {
      toStringRec(n.left, sb, depth + 1);   
      toStringRec(n.right, sb, depth + 1);
    }
  }
  
  
  
  /**
   * Iterator implementation for this binary search tree.  The elements
   * are returned in ascending order according to their natural ordering.
   */
  private class BSTIterator implements Iterator<E>
  {
    /**
     * Node to be returned by next call to next().
     */
    private Node current;
    
    /**
     * Node returned by last call to next() and available
     * for removal.  This field is null when no node is
     * available to be removed.
     */
    private Node pending;
    
    /**
     * Constructs an iterator starting at the smallest
     * element in the tree.
     */
    public BSTIterator()
    {
      // start out at smallest value
      current = root;
      if (current != null)
      {
        while (current.left != null)
        {
          current = current.left;
        }
      }
    }
    
    @Override
    public boolean hasNext()
    {
      return current != null;
    }

    @Override
    public E next()
    {
      if (!hasNext()) throw new NoSuchElementException();
      pending = current;
      current = successor(current);
      return pending.data;
    }

    @Override
    public void remove()
    {
      if (pending == null) throw new IllegalStateException();

      // Remember, current points to the successor of 
      // pending, but if pending has two children, then
      // unlinkNode(pending) will copy the successor's data 
      // into pending and delete the successor node.
      // So in this case we want to end up with current
      // pointing to the pending node.
      if (pending.left != null && pending.right != null)
      {
        current = pending;
      }
      unlinkNode(pending);
      pending = null;
    }
  }
}
